# coding=utf-8
# Copyright (C) xxx team - All Rights Reserved
# 
# @Version:   3.10.4
# @Software:  PyCharm
# @FileName:  bert_mask.py
# @CTime:     2022/7/31 19:36   
# @Author:    yhy
# @Email:     yhy@yhy.com
# @UTime:     2022/7/31 19:36
# 
# @Description:
# 
#     xxx
# 
import codecs
import logging
from typing import List, Dict, Optional

import numpy as np

logger = logging.getLogger(__name__)


def main(token_ids):
    """对输入进行随机mask
    """
    VOCAB_SIZE = 12000
    MASK = 100
    ignore = -100
    rands = np.random.random(len(token_ids))
    source, target = [], []
    for r, t in zip(rands, token_ids):
        if r < 0.15 * 0.8:
            source.append(MASK)
            target.append(t)
        elif r < 0.15 * 0.9:
            source.append(t)
            target.append(t)
        elif r < 0.15:
            source.append(np.random.choice(VOCAB_SIZE - 1) + 1)
            target.append(t)
        else:
            source.append(t)
            target.append(ignore)
    return source, target



if __name__ == '__main__':
    for _ in range(10):
        source, target = main([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
        # loss = loss_fn(torch.tensor(source), torch.Tensor(target))
        print(source, target, '\n', sep='\n')
